import json
import os
import random
from datetime import datetime
from typing import Dict, Any, Optional

class UDMDataLoader:
    """Data loader for UDM subscription information"""
    
    def __init__(self, data_dir: str = "data"):
        self.data_dir = data_dir
        self._cache = {}
        self._load_predefined_data()
    
    def _load_predefined_data(self):
        """Load all predefined MSISDN data from JSON files"""
        if not os.path.exists(self.data_dir):
            return
        
        for filename in os.listdir(self.data_dir):
            if filename.endswith('.json'):
                msisdn = filename.replace('.json', '')
                filepath = os.path.join(self.data_dir, filename)
                try:
                    with open(filepath, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        self._cache[msisdn] = data
                except Exception as e:
                    print(f"Warning: Could not load {filepath}: {e}")
    
    def get_msisdn_data(self, msisdn: str) -> Dict[str, Any]:
        """Get UDM data for a specific MSISDN"""
        # Check if we have predefined data
        if msisdn in self._cache:
            return self._cache[msisdn]
        
        # Generate random demo data for unknown MSISDNs
        return self._generate_random_data(msisdn)
    
    def _generate_random_data(self, msisdn: str) -> Dict[str, Any]:
        """Generate random demo data for unknown MSISDNs"""
        imsi = f"46000405263{random.randint(5000, 9999)}"
        imei = f"862549060127{random.randint(140, 999)}"
        
        return {
            "regUdm": f"UDM{random.randint(10, 99):03d}",
            "hlrSn": None,
            "imsi": imsi,
            "msisdn": msisdn,
            "gprsLock": random.choice(["TRUE", "FALSE"]),
            "epsLock": random.choice(["TRUE", "FALSE"]),
            "telephonyTs11": "Telephony (TS11)",
            "emgcCallTs12": "Emergency Call (TS12)",
            "smMtTs21": "Short Message MT_PP (TS21)",
            "smMoTs22": "Short Message MO_PP (TS22)",
            "generalDataCdsBs30": "General-DataCDS  (BS30)",
            "dftCallTs11": "Data CDA-9600bps (BS26)",
            "cfu": random.choice(["PROV", "NOTPROV"]),
            "cfb": random.choice(["PROV", "NOTPROV"]),
            "cfnry": random.choice(["PROV", "NOTPROV"]),
            "cfnrc": random.choice(["PROV", "NOTPROV"]),
            "cfd": random.choice(["PROV", "NOTPROV"]),
            "baoc": random.choice(["PROV", "NOTPROV"]),
            "boic": random.choice(["PROV", "NOTPROV"]),
            "cw": random.choice(["PROV", "NOTPROV"]),
            "hold": random.choice(["PROV", "NOTPROV"]),
            "gprsChargeGlobal": random.choice(["PREPAID", "POSTPAID"]),
            "epsChargeGlobal": random.choice(["PREPAID", "POSTPAID"]),
            "epsUeAmbrUl": str(random.randint(100000000, 2000000000)),
            "epsUeAmbrDl": str(random.randint(1000000000, 4294967295)),
            "psPdpChargeGlobal": random.choice(["PREPAID", "POSTPAID"]),
            "psCntxId": ["1", "2", "3"],
            "psApnQosTplId": ["540", "541", "65534"],
            "psApn": ["CMWAP", "CMNET", "*"],
            "psMbrUl": ["8192K", "8192K", "2048K"],
            "psMbrDl": ["8192K", "8192K", "2048K"],
            "epsPdnChargeGlobal": random.choice(["PREPAID", "POSTPAID"]),
            "epsDftCntxId": "2",
            "epsDftApn": "CMNET",
            "epsCntxId": ["1", "2", "6", "11", "50"],
            "epsApnQosTplId": ["507", "508", "20", "65533", "1000"],
            "epsApn": ["CMWAP", "CMNET", "IMS", "*", "DUMMY"],
            "epsQci": ["6", "6", "5", "9", "9"],
            "epsArpPri": ["6", "6", "1", "3", "3"],
            "epsApnAmbrUl": [str(random.randint(50000000, 300000000)) for _ in range(5)],
            "epsApnAmbrDl": [str(random.randint(500000000, 4000000000)) for _ in range(5)],
            "mmeHost": None,
            "mmeName": None,
            "mmeRat": "EUTRAN",
            "mmeUlTime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "mmeImei": imei,
            "mmeImeiSv": "00",
            "mmeImsVops": "SUPPORTED",
            "mmeUeSrvccCap": "UE-SRVCC-SUPPORTED",
            "ueRegAtIms": random.choice(["TRUE", "FALSE"]),
            "dynEpsCntxId": ["11", "6", "2"],
            "dynEpsApn": ["cmtd3", "ims", "cmnet"],
            "dynEpsPgwHost": [
                "topon.pgw-s5.app-hdnjihshsaegwc001bhw-08ahw011.sh.sh.node.epc.mnc000.mcc460.3gppnetwork.org",
                "topon.pgw-s5s8.shsaegw17ber.sh.sh.node.epc.mnc000.mcc460.3gppnetwork.org",
                "topon.pgw-s5s8.shsaegw17ber.sh.sh.node.epc.mnc000.mcc460.3gppnetwork.org"
            ],
            "dynEpsWildCardFlag": ["TRUE", "FALSE", "FALSE"],
            "dynEpsArd": "NOTPROV",
            "vlrNum": f"8613440{random.randint(490, 499)}",
            "csUplStatus": "Normal",
            "sgsnNum": None,
            "sgsnAddr": None,
            "ipSmGw": None,
            "ngsUeAmbrUl": f"{random.randint(2, 10)} Gbps",
            "ngsUeAmbrDl": f"{random.randint(3, 15)} Gbps",
            "ngsSNssai": ["1-ffffff"] * 13,
            "ngsDefSNssai": ["TRUE"] * 13,
            "ngsDnn": ["*", "bsjyg.sh", "nyu.sh", "cmnet", "shangda.sh", "cmwap", "lixin.sh", "shnu.sh", "fudan.sh", "sjtu.sh", "ims", "tongji.sh", "sufe.sh"],
            "ngsDefault": ["FALSE", "FALSE", "FALSE", "TRUE", "FALSE", "FALSE", "FALSE", "FALSE", "FALSE", "FALSE", "FALSE", "FALSE", "FALSE"],
            "ngsDnnQosTplId": None,
            "sscMode": ["SSC_MODE_1"] * 13,
            "allowedSscMode": ["SSC_MODE_2 | SSC_MODE_3"] * 13,
            "ngsSessionAmbrUl": [f"{random.randint(50000000, 200000000)}bps"] * 13,
            "ngsSessionAmbrDl": [f"{random.randint(500000000, 2000000000)}bps"] * 13,
            "ngsCc": [str(random.randint(200, 800))] * 13,
            "ngsQi": [str(random.randint(5, 8))] * 13,
            "ngsQiPri": ["6"] * 13,
            "ngsArpPri": ["6"] * 13,
            "ngsPreemptionCap": ["N"] * 13,
            "ngsPreemptionVul": ["Y"] * 13,
            "ngsImeisv": f"{imei}0",
            "ngsAmfInstId": f"61dc76bb-4a3a-2034-1119-030c0d0000{random.randint(10, 99):02d}",
            "ngsAmfName": f"AMF{random.randint(10, 99):03d}",
            "ngsAmfRegion": "华东南大区",
            "ngsAmfProvince": "上海",
            "ngsRegTime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "ngsAmfId": f"076{random.randint(10, 99):03d}",
            "ngsUrrpAmf": "REACHABLE",
            "ngsActDnn": ["cmnet", "ims"],
            "ngsActSmfSnssai": ["1", "1"],
            "ngsActSessionId": ["1", "2"],
            "ngsActSmfInstId": [f"621d03f5-4a6b-2034-1b7d-100c0d0000{random.randint(10, 99):02d}"] * 2,
            "ngsActSmfName": [f"SMF{random.randint(10, 99):03d}"] * 2,
            "ngsActSmfRegion": ["华东南大区"] * 2,
            "ngsActSmfProvince": ["上海"] * 2,
            "ngsSmfRegTime": [datetime.now().strftime("%Y-%m-%d %H:%M:%S")] * 2,
            "ngsSmfPgwFqdn": [
                "topon.pgw-s5.app-hdnnboshsaegwc014bhw-02ahw012.sh.sh.node.epc.mnc000.mcc460.3gppnetwork.org",
                "topon.pgw-s5.app-hdnnboshsaegwc014bhw-02ahw012.sh.sh.node.epc.mnc000.mcc460.3gppnetwork.org"
            ],
            "rawResult": f"Sample raw result data for MSISDN {msisdn}...",
            "ocsiSk": str(random.randint(1, 10)),
            "ocsiState": random.choice(["TRUE", "FALSE"]),
            "tcsiState": random.choice(["TRUE", "FALSE"]),
            "tcsiSk": str(random.randint(15, 25)),
            "vplmn": None
        }
    
    def get_available_msisdns(self) -> list:
        """Get list of available predefined MSISDNs"""
        return list(self._cache.keys())
    
    def add_msisdn_data(self, msisdn: str, data: Dict[str, Any]) -> bool:
        """Add new MSISDN data to the cache and save to file"""
        try:
            # Ensure data directory exists
            os.makedirs(self.data_dir, exist_ok=True)
            
            # Add to cache
            self._cache[msisdn] = data
            
            # Save to file
            filepath = os.path.join(self.data_dir, f"{msisdn}.json")
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=4, ensure_ascii=False)
            
            return True
        except Exception as e:
            print(f"Error saving MSISDN data: {e}")
            return False 